from . import db_connect
from . import baseFunction as bf
from process.daily_data import HS_MARKET as dailyClose
import json
import pandas as pd


def getMaAmount_forLine(tsCode: str, period: str) -> str:
    """
    获取指定 ts_code 和 period 的市场金额数据。

    参数:
    - ts_code: str - 股票代码
    - period: str - 数据周期（'W' 表示周，'M' 表示月）

    返回:
    - str: JSON 格式的返回数据，包含 trade_date, amount 和 float_mv
    """
    # 构造 SQL 查询语句
    query = f"SELECT trade_date, amount, float_mv FROM market_amount WHERE ts_code = '{tsCode}' ORDER BY trade_date"

    # 调用 db_connect.get_data_mp_alwaysup() 函数执行查询
    queryData = db_connect.get_data_mp_alwaysup(query)

    # 将查询结果转换为 DataFrame
    df = pd.DataFrame(queryData, columns=['trade_date', 'amount', 'float_mv'])

    # 检查 period 参数并进行相应的重采样操作
    if period in ['W', 'M']:
        # 将 trade_date 转换为 datetime 类型并设置为 DataFrame 的索引
        df['trade_date'] = pd.to_datetime(df['trade_date'])
        df.set_index('trade_date', inplace=True)

        resample_period = 'W' if period == 'W' else 'M'
        resampled_amount = df['amount'].resample(
            resample_period).mean().dropna().tail(250)
        resampled_float_mv = df['float_mv'].resample(
            resample_period).mean().dropna().tail(250)

        # 创建一个新的 DataFrame 包含重采样后的数据
        resampled_df = pd.DataFrame(
            {'amount': resampled_amount, 'float_mv': resampled_float_mv})

        # 将重采样后的 DataFrame 转换为原始字典格式
        result_data = {
            'trade_date': resampled_df.index.strftime('%Y-%m-%d').tolist(),
            'amount': resampled_df['amount'].astype(int).tolist(),
            'float_mv': resampled_df['float_mv'].astype(int).tolist()
        }
    else:
        # 处理未指定或不支持的 period 参数，默认为天数据
        result_data = {
            'trade_date': df['trade_date'].tail(250).tolist(),
            'amount': df['amount'].tail(250).tolist(),
            'float_mv': df['float_mv'].tail(250).tolist()
        }

    # 构造返回结果
    result = db_connect.R(code=200, success=True,
                          msg="Data fetched successfully", data=result_data)

    json_result = json.dumps(result, cls=bf.CustomEncoder)

    return json_result


def getCloseData(tradeDate: str) -> str:
    # 构造返回结果
    result_data = dailyClose.CloseData_HS_MARKET(tradeDate)
    result = db_connect.R(code=200, success=True,
                          msg="Data fetched successfully", data=result_data)

    json_result = json.dumps(result, cls=bf.CustomEncoder)
    return json_result

    # dict = getMaAmount('HS_MARKET','D')
    # print(dict)


def getCloseData_detail_circMv(tradeDate: str) -> str:
    # 构造返回结果
    result_data = dailyClose.CloseData_HS_MARKET(tradeDate)
    result = db_connect.R(code=200, success=True,
                          msg="Data fetched successfully", data=result_data)

    json_result = json.dumps(result, cls=bf.CustomEncoder)
    return json_result
